#pragma once

#ifndef __DRIVER_NVME_HH__
#define __DRIVER_NVME_HH__

#include <list>
#include <unordered_map>

#include "driver/abstract_interface.hh"
#include "driver/nvme/prp.hh"
#include "driver/nvme/queue.hh"
#include "msctrl/hil/nvme/controller.hh"
#include "msctrl/sim/msctrl.hh"
#include "msctrl/util/scheduler.hh"

#define QUEUE_ENTRY_ADMIN 256
#define QUEUE_ENTRY_IO 1024

namespace Standalone::Driver::NVMe {

// Copied from MsCtrl-FullSystem
struct DMAEntry {
  uint64_t addr;
  uint64_t size;
  uint8_t *buffer;
  MsCtrl::Event eid;
  uint64_t data;

  DMAEntry()
      : addr(0), size(0), buffer(nullptr), eid(MsCtrl::InvalidEventID) {}

  static void backup(std::ostream &, DMAEntry *) {}
  static DMAEntry *restore(std::istream &, MsCtrl::ObjectData &) {
    return nullptr;
  }
};

enum class InitState : uint8_t {
  None,
  Phase0,
  Phase1,
  Phase2,
  Phase3,
  Phase4,
  Phase5,
  Inited,
};

struct IOWrapper {
  uint64_t id;
  PRP *prp;

  IOWrapper(uint64_t i, PRP *p) : id(i), prp(p) {}
};

class NVMeInterface : public AbstractInterface, MsCtrl::Interface {
 private:
  MsCtrl::HIL::NVMe::Controller *controller;

  // PCI Express (for DMA throttling)
  MsCtrl::DelayFunction delayFunction;

  // DMA scheduling
  MsCtrl::Scheduler<DMAEntry *> scheduler;

  // NVMe Identify
  Event beginEvent;
  uint64_t capacity;
  uint32_t LBAsize;
  uint32_t namespaceID;

  // Queue
  uint32_t maxQueueEntries;
  bool phase;

  InitState initState;

  Queue *adminSQ;
  Queue *adminCQ;
  Queue *ioSQ;
  Queue *ioCQ;

  uint64_t preSubmitRead(DMAEntry *);
  uint64_t preSubmitWrite(DMAEntry *);
  void postDone(DMAEntry *);

  PRP *adminPRP;

  void _init0();
  void _init1();
  void _init2();
  void _init3(uint32_t);
  void _init4(uint16_t);
  void _init5(uint16_t);

  void submitCommand(uint16_t, uint8_t *);
  void callback(uint16_t, uint8_t *);

 public:
  NVMeInterface(ObjectData &, MsCtrl::MsCtrl &);
  ~NVMeInterface();

  // AbstractInterface
  void initialize(IGL::BlockIOLayer *, Event) override;
  void getSSDInfo(uint64_t &, uint32_t &) override;
  void submit(Request &) override;

  // MsCtrl::Interface
  void read(uint64_t, uint32_t, uint8_t *, MsCtrl::Event,
            uint64_t = 0) override;
  void write(uint64_t, uint32_t, uint8_t *, MsCtrl::Event,
             uint64_t = 0) override;

  void postInterrupt(uint16_t, bool) override;
  void getPCIID(uint16_t &, uint16_t &) override;

  void getStatList(std::vector<MsCtrl::Stat> &) override;
  void getStatValues(std::vector<double> &) override;
};

}  // namespace Standalone::Driver::NVMe

#endif
